from kernel_gen.llm_api import *
from kernel_gen.question_gen.ask_questions import ask_question_pipeline_multi_round
from kernel_gen.question_gen.cann_data import load_jsons_from_path
from kernel_gen.question_gen.transforms import TreeTraverse
config = APIConfigFactory.get_llm_api("qwq-official")
setup_default_sglang_backend(config)
# all_chapters = load_jsons_from_path("/home/ma-user/work/users/wty/kernel-gen/data/chapter_contents")
all_chapters = load_jsons_from_path("/home/ma-user/work/users/wty/kernel-gen/examples/ask_questions/chapters")
# ask_question_for_one_chapter_pipeline(all_chapters)

tranverse = TreeTraverse(question_depth=4,max_conv_round=20,
                         output_dir="/home/ma-user/work/users/wty/kernel-gen/examples/ask_questions/new_questions",
                         function_handle=ask_question_pipeline_multi_round)
tranverse.visit(all_chapters)